#include "nm_getproc.h"

PVOID kgetproc::FindExportedRoutineByName(PVOID DllBase, PANSI_STRING AnsiImageRoutineName )
{
	USHORT OrdinalNumber;
	PULONG NameTableBase;
	PUSHORT NameOrdinalTableBase;
	PULONG Addr;
	LONG High;
	LONG Low;
	LONG Middle;
	LONG Result;
	ULONG ExportSize;
	PVOID FunctionAddress;
	PIMAGE_EXPORT_DIRECTORY ExportDirectory;

	ExportDirectory = (PIMAGE_EXPORT_DIRECTORY) RtlImageDirectoryEntryToData (
		DllBase,
		TRUE,
		IMAGE_DIRECTORY_ENTRY_EXPORT,
		&ExportSize);

	if (ExportDirectory == NULL) {
		return NULL;
	}

	NameTableBase = (PULONG)((PCHAR)DllBase + (ULONG)ExportDirectory->AddressOfNames);

	NameOrdinalTableBase = (PUSHORT)((PCHAR)DllBase + (ULONG)ExportDirectory->AddressOfNameOrdinals);


	Low = 0;
	Middle = 0;
	High = ExportDirectory->NumberOfNames - 1;

	while (High >= Low) {

		Middle = (Low + High) >> 1;

		Result = strcmp (AnsiImageRoutineName->Buffer,
			(PCHAR)DllBase + NameTableBase[Middle]);

		if (Result < 0) {
			High = Middle - 1;
		}
		else if (Result > 0) {
			Low = Middle + 1;
		}
		else {
			break;
		}
	}

	if (High < Low) {
		return NULL;
	}

	OrdinalNumber = NameOrdinalTableBase[Middle];

	if ((ULONG)OrdinalNumber >= ExportDirectory->NumberOfFunctions) {
		return NULL;
	}

	Addr = (PULONG)((PCHAR)DllBase + (ULONG)ExportDirectory->AddressOfFunctions);

	FunctionAddress = (PVOID)((PCHAR)DllBase + Addr[OrdinalNumber]);

	return FunctionAddress;
}

PVOID kgetproc::GetSystemRoutineAddress(PUNICODE_STRING ModuleName, PUNICODE_STRING SystemRoutineName)
{
	NTSTATUS Status;
	PKLDR_DATA_TABLE_ENTRY DataTableEntry;
	ANSI_STRING AnsiString;
	PLIST_ENTRY NextEntry;
	PVOID FunctionAddress;
	const LARGE_INTEGER ShortTime = {(ULONG)(-10 * 1000 * 10), -1};

	FunctionAddress = NULL;

	do {
		Status = RtlUnicodeStringToAnsiString (&AnsiString,
			SystemRoutineName,
			TRUE);

		if (NT_SUCCESS (Status)) {
			break;
		}

		KeDelayExecutionThread (KernelMode, FALSE, (PLARGE_INTEGER)&ShortTime);

	} while (TRUE);

	NextEntry = m_LoadedModuleList->Flink;
	while (NextEntry != m_LoadedModuleList) {

		DataTableEntry = CONTAINING_RECORD(NextEntry,
			KLDR_DATA_TABLE_ENTRY,
			InLoadOrderLinks);

		if (DataTableEntry->BaseDllName.Buffer == NULL) {

			NextEntry = NextEntry->Flink;
			continue;
		}

		if (RtlEqualUnicodeString (ModuleName,
			&DataTableEntry->BaseDllName,
			TRUE)) {

			FunctionAddress = FindExportedRoutineByName (DataTableEntry->DllBase,
					&AnsiString);

			break;
		}

		NextEntry = NextEntry->Flink;
	}


	RtlFreeAnsiString (&AnsiString);

	return FunctionAddress;
}

